#include "header.h"

struct client_t {
  int sockfd, data_sockfd;
  struct sockaddr_in addr;
  socklen_t addrlen;
  char *root;
};


int init_data_stream(struct client_t *client) {
  int data_port;

  client->data_sockfd = -1;
  if (recv_buf(client->sockfd, &data_port, 4) <= 0) {
    return -1;
  }

  data_port = ntohl(data_port);

  if ((client->data_sockfd = socket(AF_INET, SOCK_STREAM, 0)) == -1) {
    return -1;
  }

  client->addr.sin_port = htons(data_port);
  if (connect(client->data_sockfd, &(client->addr), client->addrlen) == -1) {
    close(client->data_sockfd);
    client->data_sockfd = -1;
    return -1;
  }
  return 0;
}


int handle_get(struct client_t *client, char *path) {
  char *buffer = malloc(BUFFER_SIZE);
  int pathlen = strlen(path), len;
  FILE *file;

  strcpy(buffer, path);
  if (recv_buf(client->sockfd, &len, 4) <= 0) {
    return -1;
  }
  len = ntohl(len);
  if (pathlen + len > BUFFER_SIZE ||
      recv_buf(client->sockfd, buffer + pathlen, len) <= 0) {
    return -1;
  }
  buffer[pathlen + len] = 0;
  file = fopen(buffer, "rb");
  if (file == 0) {
    len = -1;
  } else {
    fseek(file, 0, SEEK_END);
    len = ftell(file);
    fseek(file, 0, SEEK_SET);
  }
  len = htonl(len);
  send_buf(client->sockfd, &len, 4);
  len = ntohl(len);
  if (len == -1) {
    fclose(file);
    return 0;
  }
  while (len) {
    if (len > BUFFER_SIZE) {
      fread(buffer, 1, BUFFER_SIZE, file);
      if (send_buf(client->data_sockfd, buffer, BUFFER_SIZE) == -1) {
        fclose(file);
        return -1;
      }
      len -= BUFFER_SIZE;
    } else {
      fread(buffer, 1, len, file);
      if (send_buf(client->data_sockfd, buffer, len) == -1) {
        fclose(file);
        return -1;
      }
      len = 0;
    }
  }
  fclose(file);
  return 0;
}


int handle_put(struct client_t *client, char *path) {
  char *buffer = malloc(BUFFER_SIZE);
  int pathlen = strlen(path), len, ack;
  FILE *file;

  strcpy(buffer, path);
  if (recv_buf(client->sockfd, &len, 4) <= 0) {
    return -1;
  }
  len = ntohl(len);
  if (pathlen + len > BUFFER_SIZE ||
      recv_buf(client->sockfd, buffer + pathlen, len) <= 0) {
    return -1;
  }
  buffer[pathlen + len] = 0;

  file = fopen(buffer, "wb");
  if (file == 0) {
    ack = -1;
  } else
    ack = 0;
  ack = htonl(ack);

  if (send_buf(client->sockfd, &ack, 4) == -1 || ack == -1) {
    fclose(file);
    return -1;
  }

  if (recv_buf(client->sockfd, &len, 4) <= 0) {
    fclose(file);
    return -1;
  }
  len = ntohl(len);

  while (len) {
    if (len > BUFFER_SIZE) {
      if (recv_buf(client->data_sockfd, buffer, BUFFER_SIZE) <= 0) {
        fclose(file);
        return -1;
      }
      fwrite(buffer, 1, BUFFER_SIZE, file);
      len -= BUFFER_SIZE;
    } else {
      if (recv_buf(client->data_sockfd, buffer, len) <= 0) {
        fclose(file);
        return -1;
      }
      fwrite(buffer, 1, len, file);
      len = 0;
    }
  }
  fclose(file);
  return 0;
}



int handle_pwd(struct client_t *client, char *path) {
  int len = htonl(strlen(path) - strlen(client->root));
  if (send_buf(client->sockfd, &len, 4) == -1 ||
      send_buf(client->sockfd, path + strlen(client->root), ntohl(len)) == -1) {
    return -1;
  }
  return 0;
}



int handle_dir(struct client_t *client, char *path) {
  DIR *dir = opendir(path);
  struct dirent *dp;
  char *buf = malloc(BUFFER_SIZE);
  int len = 0, cap = BUFFER_SIZE;

  if (dir == 0) {
    return -1;
  }
  while ((dp = readdir(dir)) != 0) {
    while (strlen(dp->d_name) + len + 10 >= cap) {
      realloc(buf, cap *= 2);
    }
    strcpy(buf + len, dp->d_name);
    len += strlen(dp->d_name);
    if (dp->d_type == DT_DIR)
      buf[len++] = '/';
    buf[len++] = '\n';
  }
  closedir(dir);

  len = htonl(len);
  if (send_buf(client->sockfd, &len, 4) == -1 ||
      send_buf(client->sockfd, buf, ntohl(len)) == -1) {
    free(buf);
    return -1;
  }
  free(buf);
  return 0;
}


int handle_cd(struct client_t *client, char *path) {
  char *buf = malloc(BUFFER_SIZE);
  char *buf2 = malloc(BUFFER_SIZE);
  int len, tmp;

  if (recv_buf(client->sockfd, &len, 4) <= 0 ||
      strlen(path) + ntohl(len) >= BUFFER_SIZE ||
      recv_buf(client->sockfd, buf, ntohl(len)) == -1) {
    free(buf);
    free(buf2);
    return -1;
  }
  buf[ntohl(len)] = 0;
  if (buf[0] == '/') {
    strcpy(buf2, client->root);
  } else {
    strcpy(buf2, path);
  }
  strcpy(buf2 + strlen(buf2), buf);
  realpath(buf2, buf);
  if (strlen(buf) >= strlen(client->root) &&
      memcmp(buf, client->root, strlen(client->root)) == 0) {
    strcpy(path, buf);
    tmp = strlen(path);
    path[tmp] = '/';
    path[tmp + 1] = 0;
    free(buf);
    free(buf2);
  }

  len = htonl(strlen(path) - strlen(client->root));
  if (send_buf(client->sockfd, &len, 4) == -1 ||
      send_buf(client->sockfd, path + strlen(client->root), ntohl(len)) == -1) {
    return -1;
  }
  return 0;
}


void* client_handler(void *arg) {
  struct client_t *client = arg;
  char *path = malloc(BUFFER_SIZE), op;
  int tmp;

  strcpy(path, client->root);
  tmp = strlen(path);
  path[tmp] = '/';
  path[tmp + 1] = 0;

  sync_printf("%s: login\n", inet_ntoa(client->addr.sin_addr));
  if (init_data_stream(client) == -1) {
    sync_printf("%s: Error when init_data_stream\n",
                inet_ntoa(client->addr.sin_addr));
    goto quit;
  } else {
    sync_printf("%s: init_data_stream\n",
                inet_ntoa(client->addr.sin_addr));
  }

  while (1) {
    tmp = recv_buf(client->sockfd, &op, 1);
    if (tmp == -1) {
      sync_printf("%s: Error: recv command\n",
                  inet_ntoa(client->addr.sin_addr));
      goto quit;
    }
    if (tmp == 0) {
      sync_printf("%s: quit\n",
                  inet_ntoa(client->addr.sin_addr));
      goto quit;
    }
    if (op == 1) {
      sync_printf("%s: start get\n",
                  inet_ntoa(client->addr.sin_addr));
      if (handle_get(client, path) == -1) {
        sync_printf("%s: Error: get\n",
                    inet_ntoa(client->addr.sin_addr));
        goto quit;
      }
      sync_printf("%s: get finish\n",
                  inet_ntoa(client->addr.sin_addr));
    } else if (op == 2) {
      sync_printf("%s: start put\n",
                  inet_ntoa(client->addr.sin_addr));
      if (handle_put(client, path) == -1) {
        sync_printf("%s: Error: put\n",
                    inet_ntoa(client->addr.sin_addr));
        goto quit;
      }
      sync_printf("%s: put finish\n",
                  inet_ntoa(client->addr.sin_addr));
    } else if (op == 3) {
      if (handle_pwd(client, path) == -1) {
        sync_printf("%s: Error: pwd\n",
                    inet_ntoa(client->addr.sin_addr));
        goto quit;
      } else {
        sync_printf("%s: pwd\n",
                    inet_ntoa(client->addr.sin_addr));
      }
    } else if (op == 4) {
      if (handle_dir(client, path) == -1) {
        sync_printf("%s: Error: dir\n",
        inet_ntoa(client->addr.sin_addr));
        goto quit;
      }
      sync_printf("%s: dir\n",
                  inet_ntoa(client->addr.sin_addr));
    } else if (op == 5) {
      if (handle_cd(client, path) == -1) {
        sync_printf("%s: Error: cd\n",
                    inet_ntoa(client->addr.sin_addr));
        goto quit;
      }
      sync_printf("%s: cd\n",
                  inet_ntoa(client->addr.sin_addr));

    } else {
      sync_printf("%s: unkown command\n",
                  inet_ntoa(client->addr.sin_addr));
    }
  }

quit:
  close(client->sockfd);
  close(client->data_sockfd);
  free(client);
  free(path);
  return 0;
}



int init_network() {
  int sockfd;
  struct sockaddr_in addr;
  sockfd = socket(AF_INET, SOCK_STREAM, 0);
  if (sockfd == -1) return -1;

  memset(&addr, 0, sizeof(addr));
  addr.sin_family = AF_INET;
  addr.sin_port = htons(SERVER_RECEIVE_PORT);

  if (bind(sockfd, &addr, sizeof(addr)) == -1) {
    close(sockfd);
    return -1;
  }

  if (listen(sockfd, 16) == -1) {
    close(sockfd);
    return -1;
  }
  return sockfd;
}



void main_loop(int sockfd, char *root) {
  pthread_t *thread;
  struct client_t *client;
  while (1) {
    client = malloc(sizeof(struct client_t));
    client->root = root;
    client->sockfd = accept(sockfd, &client->addr, &client->addrlen);
    if (client->sockfd == -1) {
      free(client);
      continue;
    }
    thread = malloc(sizeof(pthread_t));
    if (pthread_create(thread, 0, &client_handler, client) != 0) {
      free(client);
      continue;
    }
  }
}



int main(int argc, char *argv[]) {
  int sockfd;
  char *root = 0;

  if (argc == 2) {
    root = malloc(BUFFER_SIZE);
    realpath(argv[1], root);
  }
  if (argc != 2 || root == 0) {
    sync_printf("Usage: ./ftp_server.out [root]\n");
    sync_printf("Example: ./ftp_server.out ~/FTP\n");
    return 0;
  }

  sockfd = init_network();
  if (sockfd != -1) {
    sync_printf("Server start\n");
  } else {
    sync_printf("Fail: start server\n");
    return 0;
  }
  main_loop(sockfd, root);

  return 0;
}
